import ast
import json
import javalang
import esprima
import argparse
import os
import pickle
import time
from pathlib import Path
import subprocess
import torch.multiprocessing as mp
import os.path as osp
from datasets import load_dataset, load_from_disk
import sys
import pandas as pd
import networkx as nx
from collections import defaultdict, Counter
import re

sys.path.insert(0, "/data_ext/ref_code/LocAgent")

VERSION = "v2.3"
NODE_TYPE_DIRECTORY = "directory"
NODE_TYPE_FILE = "file"
NODE_TYPE_CLASS = "class"
NODE_TYPE_FUNCTION = "function"

EDGE_TYPE_CONTAINS = "contains"
EDGE_TYPE_INHERITS = "inherits"
EDGE_TYPE_INVOKES = "invokes"
EDGE_TYPE_IMPORTS = "imports"

VALID_NODE_TYPES = [
    NODE_TYPE_DIRECTORY,
    NODE_TYPE_FILE,
    NODE_TYPE_CLASS,
    NODE_TYPE_FUNCTION,
]
VALID_EDGE_TYPES = [
    EDGE_TYPE_CONTAINS,
    EDGE_TYPE_INHERITS,
    EDGE_TYPE_INVOKES,
    EDGE_TYPE_IMPORTS,
]

SKIP_DIRS = [".git", ".github", "target", "build", ".idea", ".vscode"]


def is_skip_dir(dirname):
    for skip_dir in SKIP_DIRS:
        if skip_dir in dirname:
            return True
    return False


class JavaAnalyzer:
    def __init__(self, filepath, source_code):
        self.filepath = filepath
        self.source_code = source_code
        self.source_lines = source_code.split("\n")
        self.nodes = []
        self.package_name = ""
        self.imports = []
        self.inheritance_info = []

    def analyze(self, tree):
        """分析Java语法树并提取节点"""
        if hasattr(tree, "package") and tree.package:
            self.package_name = tree.package.name

        # 提取import信息
        if hasattr(tree, "imports"):
            for import_decl in tree.imports:
                import_info = {
                    "path": import_decl.path,
                    "static": getattr(import_decl, "static", False),
                    "wildcard": getattr(import_decl, "wildcard", False),
                }
                self.imports.append(import_info)
                print(f"Found import: {import_info}")  # 调试信息

        # 分析类型声明
        if hasattr(tree, "types"):
            for type_decl in tree.types:
                self._visit_type(type_decl)

        return self.nodes

    def _get_line_number(self, node):
        """获取节点的行号"""
        try:
            if hasattr(node, "position") and node.position:
                return node.position.line
        except:
            pass
        return 1

    def _get_end_line_number(self, node, start_line):
        """估算节点的结束行号"""
        try:
            end_line = start_line + 10
            if hasattr(node, "body") and node.body:
                max_child_line = start_line
                for child in node.body:
                    child_line = self._get_line_number(child)
                    if child_line > max_child_line:
                        max_child_line = child_line
                end_line = max_child_line + 5
            return min(end_line, len(self.source_lines))
        except:
            return start_line + 10

    def _extract_code_block(self, start_line, end_line):
        """提取指定行范围的代码"""
        if start_line <= 0 or start_line > len(self.source_lines):
            return ""

        start_idx = start_line - 1
        end_idx = min(end_line, len(self.source_lines))

        return "\n".join(self.source_lines[start_idx:end_idx])

    def _visit_type(self, node):
        """访问类型声明节点"""
        if isinstance(node, javalang.tree.ClassDeclaration):
            self._visit_class(node)
        elif isinstance(node, javalang.tree.InterfaceDeclaration):
            self._visit_interface(node)
        elif isinstance(node, javalang.tree.EnumDeclaration):
            self._visit_enum(node)

    def _visit_class(self, node):
        """访问类声明"""
        class_name = node.name
        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node, start_line)
        code_block = self._extract_code_block(start_line, end_line)

        # 处理继承关系
        extends_class = None
        implements_interfaces = []

        if hasattr(node, "extends") and node.extends:
            extends_class = node.extends.name

        if hasattr(node, "implements") and node.implements:
            implements_interfaces = [impl.name for impl in node.implements]

        class_node = {
            "name": class_name,
            "type": NODE_TYPE_CLASS,
            "code": code_block,
            "start_line": start_line,
            "end_line": end_line,
            "extends": extends_class,
            "implements": implements_interfaces,
            "modifiers": [str(mod) for mod in node.modifiers] if node.modifiers else [],
        }
        self.nodes.append(class_node)

        # 记录继承信息
        if extends_class:
            self.inheritance_info.append((class_name, extends_class, "extends"))
        for interface in implements_interfaces:
            self.inheritance_info.append((class_name, interface, "implements"))

        # 处理类成员
        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, class_name)
                elif isinstance(member, javalang.tree.ConstructorDeclaration):
                    self._visit_constructor(member, class_name)

    def _visit_interface(self, node):
        """访问接口声明"""
        interface_name = node.name
        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node, start_line)
        code_block = self._extract_code_block(start_line, end_line)

        # 接口可以extends其他接口
        extends_interfaces = []
        if hasattr(node, "extends") and node.extends:
            extends_interfaces = [ext.name for ext in node.extends]

        interface_node = {
            "name": interface_name,
            "type": NODE_TYPE_CLASS,
            "code": code_block,
            "start_line": start_line,
            "end_line": end_line,
            "extends": extends_interfaces,
            "is_interface": True,
            "modifiers": [str(mod) for mod in node.modifiers] if node.modifiers else [],
        }
        self.nodes.append(interface_node)

        # 记录继承信息
        for parent_interface in extends_interfaces:
            self.inheritance_info.append((interface_name, parent_interface, "extends"))

        # 处理接口方法
        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, interface_name)

    def _visit_enum(self, node):
        """访问枚举声明"""
        enum_name = node.name
        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node, start_line)
        code_block = self._extract_code_block(start_line, end_line)

        enum_node = {
            "name": enum_name,
            "type": NODE_TYPE_CLASS,
            "code": code_block,
            "start_line": start_line,
            "end_line": end_line,
            "is_enum": True,
            "modifiers": [str(mod) for mod in node.modifiers] if node.modifiers else [],
        }
        self.nodes.append(enum_node)

    def _visit_method(self, node, parent_class):
        """访问方法声明"""
        method_name = node.name
        full_method_name = f"{parent_class}.{method_name}"

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node, start_line)
        code_block = self._extract_code_block(start_line, end_line)

        # 获取参数信息
        parameters = []
        if hasattr(node, "parameters") and node.parameters:
            for param in node.parameters:
                parameters.append(
                    {
                        "name": param.name,
                        "type": str(param.type) if param.type else "unknown",
                    }
                )

        # 提取方法调用
        method_calls = self._extract_method_calls_from_node(node)

        method_node = {
            "name": full_method_name,
            "type": NODE_TYPE_FUNCTION,
            "code": code_block,
            "start_line": start_line,
            "end_line": end_line,
            "parameters": parameters,
            "return_type": str(node.return_type) if node.return_type else "void",
            "modifiers": [str(mod) for mod in node.modifiers] if node.modifiers else [],
            "method_calls": method_calls,
        }
        self.nodes.append(method_node)

    def _visit_constructor(self, node, parent_class):
        """访问构造函数声明"""
        constructor_name = node.name
        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node, start_line)
        code_block = self._extract_code_block(start_line, end_line)

        parameters = []
        if hasattr(node, "parameters") and node.parameters:
            for param in node.parameters:
                parameters.append(
                    {
                        "name": param.name,
                        "type": str(param.type) if param.type else "unknown",
                    }
                )

        # 提取方法调用
        method_calls = self._extract_method_calls_from_node(node)

        constructor_node = {
            "name": constructor_name,
            "type": NODE_TYPE_FUNCTION,
            "code": code_block,
            "start_line": start_line,
            "end_line": end_line,
            "parameters": parameters,
            "is_constructor": True,
            "modifiers": [str(mod) for mod in node.modifiers] if node.modifiers else [],
            "method_calls": method_calls,
        }
        self.nodes.append(constructor_node)

    def _extract_method_calls_from_node(self, node):
        """从节点中提取方法调用"""
        calls = []

        def collect_calls(n):
            if isinstance(n, javalang.tree.MethodInvocation):
                calls.append(n.member)
            elif isinstance(n, javalang.tree.ClassCreator):
                # 构造函数调用
                if hasattr(n, "type") and hasattr(n.type, "name"):
                    calls.append(n.type.name)

            # 递归遍历子节点
            if hasattr(n, "children"):
                for child in n.children:
                    if child:
                        if isinstance(child, list):
                            for item in child:
                                if item:
                                    collect_calls(item)
                        else:
                            collect_calls(child)

        try:
            collect_calls(node)
        except Exception as e:
            print(f"Error extracting method calls: {e}")

        return calls


def analyze_file_java(filepath):
    """分析单个Java文件，返回类和方法节点"""
    try:
        with open(filepath, "r", encoding="utf-8") as file:
            code = file.read()

        tree = javalang.parse.parse(code)
        analyzer = JavaAnalyzer(filepath, code)
        nodes = analyzer.analyze(tree)
        return nodes, analyzer.imports, analyzer.inheritance_info
    except Exception as e:
        print(f"Error parsing {filepath}: {e}")
        return [], [], []


def resolve_java_import(import_path, repo_path, all_java_files):
    """解析Java import路径到文件路径"""
    print(f"Resolving import: {import_path}")

    if import_path.endswith(".*"):
        # 通配符导入，返回包目录下的所有Java文件
        package_path = import_path[:-2].replace(".", "/")
        print(f"  Package path: {package_path}")

        matching_files = []
        for java_file in all_java_files:
            file_dir = os.path.dirname(java_file)
            if file_dir.endswith(package_path):
                matching_files.append(java_file)
                print(f"  Found matching file: {java_file}")

        return matching_files
    else:
        # 具体类导入
        class_name = import_path.split(".")[-1]
        package_path = "/".join(import_path.split(".")[:-1])

        print(f"  Class name: {class_name}, Package path: {package_path}")

        # 查找匹配的文件
        for java_file in all_java_files:
            file_name = os.path.basename(java_file)
            file_dir = os.path.dirname(java_file)

            # 检查文件名是否匹配
            if file_name == f"{class_name}.java":
                # 检查包路径是否匹配
                if package_path == "" or file_dir.endswith(package_path):
                    print(f"  Found matching file: {java_file}")
                    return [java_file]

        print(f"  No matching file found for {import_path}")
        return []


def add_java_imports(root_node, imports, graph, repo_path, all_java_files):
    """添加Java导入边"""
    print(f"Adding imports for {root_node}")

    for imp in imports:
        import_path = imp["path"]
        resolved_files = resolve_java_import(import_path, repo_path, all_java_files)

        for resolved_file in resolved_files:
            if graph.has_node(resolved_file):
                graph.add_edge(
                    root_node,
                    resolved_file,
                    type=EDGE_TYPE_IMPORTS,
                    static=imp.get("static", False),
                    wildcard=imp.get("wildcard", False),
                )
                print(f"  Added import edge: {root_node} -> {resolved_file}")
            else:
                print(f"  Node not found in graph: {resolved_file}")


def get_java_inner_nodes(query_node, src_node, graph):
    """获取Java节点的内部节点"""
    inner_nodes = []
    try:
        for _, dst_node, attr in graph.edges(src_node, data=True):
            if attr["type"] == EDGE_TYPE_CONTAINS and dst_node != query_node:
                inner_nodes.append(dst_node)
                if graph.nodes[dst_node]["type"] == NODE_TYPE_CLASS:
                    inner_nodes.extend(
                        get_java_inner_nodes(query_node, dst_node, graph)
                    )
    except:
        pass
    return inner_nodes


def find_all_possible_callee_java_fixed(node, graph):
    """
    找到Java节点的所有可能被调用方
    修复版本：确保只返回已导入的函数
    """
    callee_nodes = []
    cur_node = node
    pre_node = node

    # 追踪已导入的文件和类
    imported_files = set()
    imported_classes = set()

    def find_parent(_cur_node):
        """找到包含当前节点的父节点"""
        try:
            for predecessor in graph.predecessors(_cur_node):
                edge_data = graph.get_edge_data(predecessor, _cur_node)
                if edge_data:
                    for key, attr in edge_data.items():
                        if attr["type"] == EDGE_TYPE_CONTAINS:
                            return predecessor
        except:
            pass
        return None

    def get_file_node(node):
        """获取节点所在的文件节点"""
        current = node
        while current:
            if graph.nodes[current]["type"] == NODE_TYPE_FILE:
                return current
            current = find_parent(current)
        return None

    # 首先找到当前节点所在的文件
    file_node = get_file_node(node)

    # 收集该文件的所有导入
    if file_node:
        try:
            for _, dst_node, attr in graph.edges(file_node, data=True):
                if attr["type"] == EDGE_TYPE_IMPORTS:
                    imported_files.add(dst_node)
                    print(f"Found imported file: {dst_node}")

                    # 收集导入文件中的所有类和函数
                    for _, inner_node, inner_attr in graph.edges(dst_node, data=True):
                        if inner_attr["type"] == EDGE_TYPE_CONTAINS:
                            node_type = graph.nodes[inner_node]["type"]
                            if node_type in [NODE_TYPE_CLASS, NODE_TYPE_FUNCTION]:
                                imported_classes.add(inner_node)
        except Exception as e:
            print(f"Error collecting imports: {e}")

    # 遍历节点层次结构
    while cur_node:
        # 获取当前层级的内部节点（同一文件/类内的节点）
        inner_nodes = get_java_inner_nodes(pre_node, cur_node, graph)
        callee_nodes.extend(inner_nodes)

        if graph.nodes[cur_node]["type"] == NODE_TYPE_FILE:
            # 当到达文件节点时，只添加已导入的节点
            for imported_file in imported_files:
                if graph.has_node(imported_file):
                    # 获取导入文件中的可调用节点
                    imported_inner_nodes = get_java_inner_nodes(cur_node, imported_file, graph)
                    callee_nodes.extend(imported_inner_nodes)

                    # 直接添加导入的类和函数
                    for imp_class in imported_classes:
                        if graph.has_node(imp_class):
                            # 检查是否属于当前导入的文件
                            imp_file = get_file_node(imp_class)
                            if imp_file == imported_file:
                                callee_nodes.append(imp_class)
            break

        pre_node = cur_node
        cur_node = find_parent(cur_node)

    # 去重并返回
    return list(set(callee_nodes))


def build_java_graph(repo_path, fuzzy_search=True, global_import=False):
    """构建Java项目的依赖图"""
    graph = nx.MultiDiGraph()
    file_nodes = {}
    file_analysis_results = {}  # 存储文件分析结果，避免重复解析
    all_inheritance_info = []

    print("Building directory structure...")

    # 添加根目录节点
    graph.add_node("/", type=NODE_TYPE_DIRECTORY)
    dir_stack = []
    dir_include_stack = []

    # 收集所有Java文件路径
    all_java_files = []
    for root, _, files in os.walk(repo_path):
        for file in files:
            if file.endswith(".java"):
                file_path = os.path.join(root, file)
                filename = os.path.relpath(file_path, repo_path)
                all_java_files.append(filename)

    print(f"Found {len(all_java_files)} Java files")

    for root, _, files in os.walk(repo_path):
        # 添加目录节点和边
        dirname = os.path.relpath(root, repo_path)
        if dirname == ".":
            dirname = "/"
        elif is_skip_dir(dirname):
            continue
        else:
            graph.add_node(dirname, type=NODE_TYPE_DIRECTORY)
            parent_dirname = os.path.dirname(dirname)
            if parent_dirname == "":
                parent_dirname = "/"
            graph.add_edge(parent_dirname, dirname, type=EDGE_TYPE_CONTAINS)

        # 处理目录堆栈
        while len(dir_stack) > 0 and not dirname.startswith(dir_stack[-1]):
            if not dir_include_stack[-1]:
                graph.remove_node(dir_stack[-1])
            dir_stack.pop()
            dir_include_stack.pop()

        if dirname != "/":
            dir_stack.append(dirname)
            dir_include_stack.append(False)

        dir_has_java = False
        for file in files:
            if file.endswith(".java"):
                dir_has_java = True
                try:
                    file_path = os.path.join(root, file)
                    filename = os.path.relpath(file_path, repo_path)

                    if os.path.islink(file_path):
                        continue

                    with open(file_path, "r", encoding="utf-8") as f:
                        file_content = f.read()

                    graph.add_node(filename, type=NODE_TYPE_FILE, code=file_content)
                    file_nodes[filename] = file_path

                    # 分析Java文件（只解析一次）
                    nodes, imports, inheritance_info = analyze_file_java(file_path)
                    file_analysis_results[filename] = {
                        "nodes": nodes,
                        "imports": imports,
                        "inheritance_info": inheritance_info,
                    }

                    all_inheritance_info.extend(
                        [
                            (filename, child, parent, rel_type)
                            for child, parent, rel_type in inheritance_info
                        ]
                    )

                    # 添加类和方法节点
                    for node in nodes:
                        full_name = f'{filename}:{node["name"]}'
                        graph.add_node(
                            full_name,
                            type=node["type"],
                            code=node["code"],
                            start_line=node["start_line"],
                            end_line=node["end_line"],
                            method_calls=node.get("method_calls", []),
                        )

                    # 添加包含关系边
                    graph.add_edge(dirname, filename, type=EDGE_TYPE_CONTAINS)
                    for node in nodes:
                        full_name = f'{filename}:{node["name"]}'
                        name_list = node["name"].split(".")
                        if len(name_list) == 1:
                            graph.add_edge(filename, full_name, type=EDGE_TYPE_CONTAINS)
                        else:
                            parent_name = ".".join(name_list[:-1])
                            full_parent_name = f"{filename}:{parent_name}"
                            if graph.has_node(full_parent_name):
                                graph.add_edge(
                                    full_parent_name, full_name, type=EDGE_TYPE_CONTAINS
                                )

                except (UnicodeDecodeError, Exception) as e:
                    print(f"Error processing {file_path}: {e}")
                    continue

        # 保留包含Java文件的父目录
        if dir_has_java:
            for i in range(len(dir_include_stack)):
                dir_include_stack[i] = True

    # 清理最后的目录栈
    while len(dir_stack) > 0:
        if not dir_include_stack[-1]:
            graph.remove_node(dir_stack[-1])
        dir_stack.pop()
        dir_include_stack.pop()

    print("Adding import edges...")
    # 添加导入边（使用保存的分析结果）
    import_count = 0
    for filename in file_nodes.keys():
        try:
            imports = file_analysis_results[filename]["imports"]
            if imports:
                print(f"Processing imports for {filename}: {len(imports)} imports")
                add_java_imports(filename, imports, graph, repo_path, all_java_files)
                import_count += len(imports)
        except Exception as e:
            print(f"Error processing imports for {filename}: {e}")
            continue

    print(f"Processed {import_count} import statements")

    print("Adding inheritance edges...")
    # 添加继承边
    for filename, child_class, parent_class, rel_type in all_inheritance_info:
        child_full_name = f"{filename}:{child_class}"

        # 尝试在同一文件中查找父类
        parent_full_name = f"{filename}:{parent_class}"
        if graph.has_node(parent_full_name):
            graph.add_edge(child_full_name, parent_full_name, type=EDGE_TYPE_INHERITS)
        else:
            # 在其他文件中查找父类
            for other_filename in file_nodes.keys():
                other_parent_full_name = f"{other_filename}:{parent_class}"
                if graph.has_node(other_parent_full_name):
                    graph.add_edge(
                        child_full_name, other_parent_full_name, type=EDGE_TYPE_INHERITS
                    )
                    break

    print("Adding method invocation edges...")
    # 全局名称字典（仅在global_import开启时使用）
    global_name_dict = defaultdict(list)
    if global_import:
        for node in graph.nodes():
            if ":" in node:
                node_name = node.split(":")[-1].split(".")[-1]
                global_name_dict[node_name].append(node)

    # 添加方法调用边
    method_nodes = [
        node
        for node, attr in graph.nodes(data=True)
        if attr.get("type") == NODE_TYPE_FUNCTION
    ]

    invocation_count = 0
    for node in method_nodes:
        try:
            # 获取方法调用信息
            method_calls = graph.nodes[node].get("method_calls", [])

            if not method_calls:
                continue

            # 使用修复后的函数找到所有可能的被调用方
            callee_nodes = find_all_possible_callee_java_fixed(node, graph)

            print(f"Processing invocations for {node}: {len(method_calls)} calls, {len(callee_nodes)} possible callees")

            if fuzzy_search:
                callee_name_dict = defaultdict(list)
                for callee_node in set(callee_nodes):
                    if ":" in callee_node:
                        callee_name = callee_node.split(":")[-1].split(".")[-1]
                        callee_name_dict[callee_name].append(callee_node)
            else:
                callee_name_dict = {}
                for callee_node in set(callee_nodes):
                    if ":" in callee_node:
                        callee_name = callee_node.split(":")[-1].split(".")[-1]
                        callee_name_dict[callee_name] = callee_node

            # 添加调用边
            for callee_name in set(method_calls):
                callee_node_list = callee_name_dict.get(callee_name)
                if callee_node_list:
                    if isinstance(callee_node_list, list):
                        for callee in callee_node_list:
                            graph.add_edge(node, callee, type=EDGE_TYPE_INVOKES)
                            invocation_count += 1
                            print(f"  Added invocation edge: {node} -> {callee}")
                    else:
                        graph.add_edge(node, callee_node_list, type=EDGE_TYPE_INVOKES)
                        invocation_count += 1
                        print(f"  Added invocation edge: {node} -> {callee_node_list}")
                elif global_import:
                    # 仅在开启global_import且本地找不到时，才尝试全局搜索
                    global_fuzzy_nodes = global_name_dict.get(callee_name)
                    if global_fuzzy_nodes:
                        for global_fuzzy_node in global_fuzzy_nodes:
                            # 验证该全局节点是否可访问（通过某种导入链）
                            graph.add_edge(
                                node, global_fuzzy_node, type=EDGE_TYPE_INVOKES
                            )
                            invocation_count += 1
                            print(f"  Added global invocation edge: {node} -> {global_fuzzy_node}")
        except Exception as e:
            print(f"Error processing method calls for {node}: {e}")
            continue

    print(f"Added {invocation_count} invocation edges")
    print("Graph construction completed.")
    return graph


def save_graph_to_excel(graph, output_dir):
    """保存图信息到Excel文件"""
    os.makedirs(output_dir, exist_ok=True)

    # 保存节点信息
    nodes_data = []
    for node, attributes in graph.nodes(data=True):
        node_info = {
            "node_id": node,
            "type": attributes.get("type", ""),
            "start_line": attributes.get("start_line", ""),
            "end_line": attributes.get("end_line", ""),
            "code_length": (
                len(attributes.get("code", "")) if attributes.get("code") else 0
            ),
            "has_code": bool(attributes.get("code")),
        }
        nodes_data.append(node_info)

    nodes_df = pd.DataFrame(nodes_data)
    nodes_excel_path = os.path.join(output_dir, "nodes_info.xlsx")
    nodes_df.to_excel(nodes_excel_path, index=False)
    print(f"Nodes information saved to {nodes_excel_path}")

    # 保存边信息
    edges_data = []
    for u, v, key, attributes in graph.edges(keys=True, data=True):
        edge_info = {
            "source": u,
            "target": v,
            "edge_key": key,
            "type": attributes.get("type", ""),
        }
        edges_data.append(edge_info)

    edges_df = pd.DataFrame(edges_data)
    edges_excel_path = os.path.join(output_dir, "edges_info.xlsx")
    edges_df.to_excel(edges_excel_path, index=False)
    print(f"Edges information saved to {edges_excel_path}")

    # 保存统计信息
    stats_data = {
        "Metric": ["Total Nodes", "Total Edges"]
        + [f"Nodes - {ntype}" for ntype in VALID_NODE_TYPES]
        + [f"Edges - {etype}" for etype in VALID_EDGE_TYPES],
        "Count": [len(graph.nodes()), len(graph.edges())]
        + [
            sum(1 for _, data in graph.nodes(data=True) if data.get("type") == ntype)
            for ntype in VALID_NODE_TYPES
        ]
        + [
            sum(1 for _, _, data in graph.edges(data=True) if data.get("type") == etype)
            for etype in VALID_EDGE_TYPES
        ],
    }

    stats_df = pd.DataFrame(stats_data)
    stats_excel_path = os.path.join(output_dir, "graph_statistics.xlsx")
    stats_df.to_excel(stats_excel_path, index=False)
    print(f"Graph statistics saved to {stats_excel_path}")


def list_folders(path):
    return [p.name for p in Path(path).iterdir() if p.is_dir()]


def run(rank, repo_name, repo_path, out_path, download_repo=False, instance_data=None):
    """处理单个仓库"""
    repo_basename = os.path.basename(repo_name.rstrip("/"))
    output_file = f"{osp.join(out_path, repo_basename)}.pkl"
    repo_dir = repo_name if os.path.isabs(repo_name) else osp.join(repo_path, repo_name)

    print(f"[{rank}] Start process {repo_basename}")

    try:
        if os.path.exists(repo_dir):
            # 使用新的build_java_graph函数
            graph = build_java_graph(repo_dir, fuzzy_search=True, global_import=True)

            # 保存原始pickle文件
            with open(output_file, "wb") as f:
                pickle.dump(graph, f)

            # 保存Excel文件
            excel_output_dir = f"{out_path}/{repo_basename}_analysis"
            save_graph_to_excel(graph, excel_output_dir)

            # 打印统计信息
            node_types = []
            for node, data in graph.nodes(data=True):
                node_types.append(data["type"])

            edge_types = []
            for u, v, data in graph.edges(data=True):
                edge_types.append(data["type"])

            print(f"[{rank}] Processed {repo_basename}")
            print(f"[{rank}] Node types: {Counter(node_types)}")
            print(f"[{rank}] Edge types: {Counter(edge_types)}")
            print(
                f"[{rank}] Total nodes: {len(graph.nodes())}, Total edges: {len(graph.edges())}"
            )

        else:
            print(f"[{rank}] Repository directory not found: {repo_dir}")
    except Exception as e:
        print(f"[{rank}] Error processing {repo_basename}: {e}")
        import traceback

        traceback.print_exc()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, default="pp_data")
    parser.add_argument("--split", type=str, default="test")
    parser.add_argument("--num_processes", type=int, default=1)
    parser.add_argument("--download_repo", action="store_true", default=False)
    parser.add_argument(
        "--repo_path", type=str, default="/data_ext/ref_code/LocAgent/data/data/pp_data"
    )
    parser.add_argument("--index_dir", type=str, default="index_data_fixed")
    parser.add_argument("--instance_id_path", type=str, default="")
    args = parser.parse_args()

    dataset_name = args.dataset.split("/")[-1]
    args.index_dir = f"{args.index_dir}/{dataset_name}/graph_index_{VERSION}/"
    os.makedirs(args.index_dir, exist_ok=True)

    # 测试单个仓库
    repo_name = "/data_ext/ref_code/LocAgent/data/data/pp_data/pingp-auto-test"
    start_time = time.time()

    run(0, repo_name, args.repo_path, args.index_dir, args.download_repo, None)

    end_time = time.time()
    print(f"Total Execution time = {end_time - start_time:.3f}s")